#include <acl/acl.h>
#include <atb/atb_infer.h>
#include <iostream>

int main(int argc, const char *argv[])
{
    int deviceId = 0;
    aclError status = aclrtSetDevice(deviceId);

    atb::infer::LinearParam param;
    param.transposeA = false;
    param.transposeB = true;
    param.hasBias = false;
    param.outDataType = ACL_DT_UNDEFINED;

    atb::Operation *op = nullptr;
    atb::Status st = atb::CreateOperation(param, &op);

    atb::Tensor x;
    x.desc.dtype = ACL_FLOAT16;
    x.desc.format = ACL_FORMAT_ND;
    x.desc.shape.dimNum = 2;
    x.desc.shape.dims[0] = 16;
    x.desc.shape.dims[1] = 32;
    x.dataSize = atb::Utils::GetTensorSize(x);
    status = aclrtMalloc(&x.deviceData, x.dataSize, ACL_MEM_MALLOC_HUGE_FIRST);

    atb::Tensor weight;
    weight.desc.dtype = ACL_FLOAT16;
    weight.desc.format = ACL_FORMAT_ND;
    weight.desc.shape.dimNum = 2;
    weight.desc.shape.dims[0] = 16;
    weight.desc.shape.dims[1] = 32;
    weight.dataSize = atb::Utils::GetTensorSize(weight);
    status = aclrtMalloc(&weight.deviceData, weight.dataSize, ACL_MEM_MALLOC_HUGE_FIRST);

    atb::Tensor output;
    output.desc.dtype = ACL_FLOAT16;
    output.desc.format = ACL_FORMAT_ND;
    output.desc.shape.dimNum = 2;
    output.desc.shape.dims[0] = 16;
    output.desc.shape.dims[1] = 16;
    output.dataSize = atb::Utils::GetTensorSize(output);
    status = aclrtMalloc(&output.deviceData, output.dataSize, ACL_MEM_MALLOC_HUGE_FIRST);

    atb::VariantPack variantPack;
    variantPack.inTensors = { x, weight };
    variantPack.outTensors = { output };

    atb::Context *context = nullptr;
    st = atb::CreateContext(&context);
    aclrtStream stream = nullptr;
    status = aclrtCreateStream(&stream);
    context->SetExecuteStream(stream);
    
    uint64_t workspaceSize = 0;
    st = op->Setup(variantPack, workspaceSize, context);
    void *workspace = nullptr;
    if (workspaceSize > 0) {
        status = aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    }
    st = op->Execute(variantPack, (uint8_t *)workspace, workspaceSize, context);
    std::cout << "DEMO END" << std::endl;
}

